import baostock as bs
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
import numpy as np
import os
from datetime import datetime
import pymysql
from sqlalchemy import create_engine
import warnings
import time

warnings.filterwarnings('ignore')


# 设置中文字体
def set_chinese_font():
    """设置中文字体，解决中文显示问题"""
    try:
        # 尝试使用系统中已有的中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans Fallback']
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
        print("中文字体设置成功")
    except:
        print("使用默认字体，中文可能显示为方块")


set_chinese_font()

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'port': 3306,
    'user': 'root',
    'password': "",
    'database': 'stock_db',
    'charset': 'utf8mb4'
}


def get_mysql_connection():
    """获取MySQL数据库连接"""
    try:
        connection = pymysql.connect(**DB_CONFIG)
        return connection
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None


def get_all_stocks_from_db(page_size=100):
    """分页获取所有股票代码和名称"""
    connection = get_mysql_connection()
    if connection is None:
        return []

    all_stocks = []
    page = 0

    try:
        while True:
            offset = page * page_size
            with connection.cursor() as cursor:
                sql = "SELECT code, name FROM a_stock_basic_info WHERE code NOT IN (SELECT distinct stock_code FROM stock_chart_images_new) LIMIT %s OFFSET %s"
                cursor.execute(sql, (page_size, offset))
                results = cursor.fetchall()

                if not results:
                    break

                for code, name in results:
                    all_stocks.append((code, name))

                print(f"已获取第 {page + 1} 页股票数据，共 {len(results)} 只股票")
                page += 1

                # 添加短暂延迟，避免对数据库造成压力
                time.sleep(0.1)

    except Exception as e:
        print(f"查询股票列表失败: {e}")
    finally:
        connection.close()

    return all_stocks


def save_image_path_to_db(code, chart_type, period, image_path):
    """保存图片路径到数据库"""
    connection = get_mysql_connection()
    if connection is None:
        return False

    try:
        with connection.cursor() as cursor:
            # 先检查记录是否存在
            check_sql = """
            SELECT id FROM stock_chart_images_new 
            WHERE stock_code = %s AND chart_type = %s AND period = %s
            """
            cursor.execute(check_sql, (code, chart_type, period))
            result = cursor.fetchone()

            if result:
                # 更新现有记录
                update_sql = """
                UPDATE stock_chart_images_new 
                SET image_path = %s, update_time = NOW() 
                WHERE stock_code = %s AND chart_type = %s AND period = %s
                """
                cursor.execute(update_sql, (image_path, code, chart_type, period))
            else:
                # 插入新记录
                insert_sql = """
                INSERT INTO stock_chart_images_new 
                (stock_code, chart_type, period, image_path, create_time, update_time)
                VALUES (%s, %s, %s, %s, NOW(), NOW())
                """
                cursor.execute(insert_sql, (code, chart_type, period, image_path))

            connection.commit()
            return True
    except Exception as e:
        print(f"保存图片路径到数据库失败: {e}")
        connection.rollback()
        return False
    finally:
        connection.close()


def get_stock_kline_baostock(code, frequency="m", start_date="2010-01-01", end_date="2023-12-31"):
    """获取股票K线数据
    frequency: m-月, w-周, d-日 (baostock不支持直接获取季K线和年K线)
    """
    # 登陆系统
    lg = bs.login()
    if lg.error_code != '0':
        print(f'Baostock登录失败: {lg.error_msg}')
        return None

    # 获取股票代码
    if code.startswith('6'):
        stock_code = f"sh.{code}"
    else:
        stock_code = f"sz.{code}"

    # 获取K线数据
    rs = bs.query_history_k_data_plus(
        stock_code,
        "date,code,open,high,low,close,volume,amount,turn",
        start_date=start_date,
        end_date=end_date,
        frequency=frequency,
        adjustflag="3"  # 不复权
    )

    if rs.error_code != '0':
        print(f"查询错误: {rs.error_msg}")
        bs.logout()
        return None

    # 转换数据格式
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())

    # 如果没有数据，直接返回None
    if not data_list:
        bs.logout()
        return None

    df = pd.DataFrame(data_list, columns=rs.fields)

    # 转换数据类型
    numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'amount', 'turn']
    for col in numeric_columns:
        df[col] = pd.to_numeric(df[col], errors='coerce')

    df['date'] = pd.to_datetime(df['date'])
    df = df.set_index('date').sort_index()

    # 登出系统
    bs.logout()

    return df


def resample_to_quarterly(data):
    """将月K线数据重采样为季K线"""
    if data is None or len(data) == 0:
        return None

    # 重采样为季度数据（取每季度最后一个月的数据）
    quarterly_data = data.resample('Q').last()

    # 重新计算开盘、最高、最低价
    quarterly_open = data['open'].resample('Q').first()
    quarterly_high = data['high'].resample('Q').max()
    quarterly_low = data['low'].resample('Q').min()
    quarterly_volume = data['volume'].resample('Q').sum()
    quarterly_amount = data['amount'].resample('Q').sum()

    quarterly_data['open'] = quarterly_open
    quarterly_data['high'] = quarterly_high
    quarterly_data['low'] = quarterly_low
    quarterly_data['volume'] = quarterly_volume
    quarterly_data['amount'] = quarterly_amount

    return quarterly_data


def resample_to_yearly(data):
    """将月K线数据重采样为年K线"""
    if data is None or len(data) == 0:
        return None

    # 重采样为年度数据（取每年最后一个月的数据）
    yearly_data = data.resample('Y').last()

    # 重新计算开盘、最高、最低价
    yearly_open = data['open'].resample('Y').first()
    yearly_high = data['high'].resample('Y').max()
    yearly_low = data['low'].resample('Y').min()
    yearly_volume = data['volume'].resample('Y').sum()
    yearly_amount = data['amount'].resample('Y').sum()

    yearly_data['open'] = yearly_open
    yearly_data['high'] = yearly_high
    yearly_data['low'] = yearly_low
    yearly_data['volume'] = yearly_volume
    yearly_data['amount'] = yearly_amount

    return yearly_data


def get_all_kline_data(code, start_date="2010-01-01", end_date="2023-12-31"):
    """获取所有周期的K线数据：年、季、月"""
    print(f"开始获取股票 {code} 的K线数据...")

    # 获取月K线（baostock支持）
    monthly_data = get_stock_kline_baostock(code, "m", start_date, end_date)

    if monthly_data is None or len(monthly_data) == 0:
        print(f"股票 {code} 月K线数据获取失败，无法生成其他周期数据")
        return None

    # 基于月K线生成季K线和年K线
    quarterly_data = resample_to_quarterly(monthly_data)
    yearly_data = resample_to_yearly(monthly_data)

    return {
        'monthly': monthly_data,
        'quarterly': quarterly_data,
        'yearly': yearly_data
    }


def calculate_technical_indicators(data):
    """计算技术指标"""
    if data is None or len(data) == 0:
        return data

    # 移动平均线
    window_5 = min(5, len(data))
    window_10 = min(10, len(data))

    data['MA5'] = data['close'].rolling(window=window_5).mean()
    data['MA10'] = data['close'].rolling(window=window_10).mean()

    # 对于数据点较多的，计算MA20
    if len(data) > 20:
        data['MA20'] = data['close'].rolling(window=20).mean()

    # 相对强弱指数 (RSI)
    if len(data) > 14:
        delta = data['close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
        rs = gain / loss
        data['RSI'] = 100 - (100 / (1 + rs))

    return data


def create_comparison_chart_jpg(data_dict, code, stock_name="", save_dir="gupiaoimage"):
    """创建对比图表并保存为JPG - 优化版本"""
    # 创建股票专属目录
    stock_dir = os.path.join(save_dir, code)
    if not os.path.exists(stock_dir):
        os.makedirs(stock_dir)

    # 减小图片尺寸和DPI
    fig, axes = plt.subplots(3, 1, figsize=(12, 10))  # 从 (16, 14) 减小到 (12, 10)

    periods = ['monthly', 'quarterly', 'yearly']
    titles = ['月K线', '季K线', '年K线']
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1']

    for i, (period, title, color) in enumerate(zip(periods, titles, colors)):
        data = data_dict.get(period)
        if data is not None and len(data) > 0:
            ax = axes[i]

            # 计算技术指标
            data_with_indicators = calculate_technical_indicators(data.copy())

            # 绘制收盘价线 - 减小线宽和标记尺寸
            ax.plot(data_with_indicators.index, data_with_indicators['close'],
                    linewidth=2, color=color, label=f'{title}收盘价', marker='o', markersize=2)  # 线宽从3减小到2，标记从3减小到2

            # 添加移动平均线 - 减小线宽
            if 'MA5' in data_with_indicators.columns:
                ax.plot(data_with_indicators.index, data_with_indicators['MA5'],
                        '--', color='orange', alpha=0.7, label='MA5', linewidth=1)  # 线宽从1.5减小到1
            if 'MA10' in data_with_indicators.columns:
                ax.plot(data_with_indicators.index, data_with_indicators['MA10'],
                        '--', color='purple', alpha=0.7, label='MA10', linewidth=1)  # 线宽从1.5减小到1

            ax.set_title(f'{stock_name} {title}走势图', fontsize=14, fontweight='bold', pad=15)  # 字体从16减小到14
            ax.set_ylabel('价格 (元)', fontsize=10)  # 字体从12减小到10
            ax.legend(fontsize=8)  # 字体从10减小到8
            ax.grid(True, alpha=0.3)

            # 添加统计信息 - 减小字体
            last_price = data_with_indicators['close'].iloc[-1]
            max_price = data_with_indicators['high'].max()
            min_price = data_with_indicators['low'].min()
            change_pct = ((data_with_indicators['close'].iloc[-1] - data_with_indicators['close'].iloc[0]) /
                          data_with_indicators['close'].iloc[0] * 100)

            stats_text = f'最新: {last_price:.2f}\n最高: {max_price:.2f}\n最低: {min_price:.2f}\n总涨幅: {change_pct:.1f}%'
            ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, verticalalignment='top',
                    bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.8),
                    fontsize=8, linespacing=1.5)  # 字体从10减小到8

    plt.tight_layout(pad=2.0)  # 从3.0减小到2.0

    # 保存为JPG - 降低DPI和质量
    jpg_filename = os.path.join(stock_dir, f'{code}_comparison_chart.jpg')
    plt.savefig(jpg_filename, dpi=150, bbox_inches='tight', format='jpg')  # DPI从300降低到150，添加压缩优化，质量85%
    print(f"对比图表已保存为: {jpg_filename}")
    plt.close()  # 关闭图表以释放内存

    # 获取文件大小
    file_size = os.path.getsize(jpg_filename) / 1024  # KB
    print(f"文件大小: {file_size:.1f} KB")

    # 保存图片路径到数据库
    save_image_path_to_db(code, 'comparison', 'all', jpg_filename)

    return jpg_filename


def create_individual_chart_jpg(data_dict, code, stock_name="", save_dir="gupiaoimagenew"):
    """为年K线和季K线创建单独的JPG图表 - 优化版本"""
    # 创建股票专属目录
    stock_dir = os.path.join(save_dir, code)
    if not os.path.exists(stock_dir):
        os.makedirs(stock_dir)

    jpg_files = []

    # 只处理季K线和年K线
    target_freqs = ['quarterly', 'yearly']

    for freq in target_freqs:
        data = data_dict.get(freq)
        if data is not None and len(data) > 0:
            # 计算技术指标
            data_with_indicators = calculate_technical_indicators(data.copy())

            # 创建图表 - 减小尺寸
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))  # 从 (15, 10) 减小到 (12, 8)

            freq_name = {'quarterly': '季', 'yearly': '年'}.get(freq, freq)
            color = {'quarterly': '#4ECDC4', 'yearly': '#45B7D1'}.get(freq, '#666666')

            # 主图：价格走势 - 减小线宽和标记
            ax1.plot(data_with_indicators.index, data_with_indicators['close'],
                     linewidth=2, color=color, label='收盘价', marker='o', markersize=3)  # 线宽从3减小到2，标记从4减小到3

            # 添加移动平均线 - 减小线宽
            if 'MA5' in data_with_indicators.columns:
                ax1.plot(data_with_indicators.index, data_with_indicators['MA5'],
                         '--', color='orange', alpha=0.8, label='MA5', linewidth=1.5)  # 线宽从2减小到1.5
            if 'MA10' in data_with_indicators.columns:
                ax1.plot(data_with_indicators.index, data_with_indicators['MA10'],
                         '--', color='purple', alpha=0.8, label='MA10', linewidth=1.5)  # 线宽从2减小到1.5

            ax1.set_title(f'{stock_name} {freq_name}K线走势图', fontsize=14, fontweight='bold', pad=15)  # 字体从18减小到14
            ax1.set_ylabel('价格 (元)', fontsize=12)  # 字体从14减小到12
            ax1.legend(fontsize=10)  # 字体从12减小到10
            ax1.grid(True, alpha=0.3)

            # 添加价格统计信息 - 减小字体
            price_stats = (
                f'最新价: {data_with_indicators["close"].iloc[-1]:.2f}\n'
                f'最高价: {data_with_indicators["high"].max():.2f}\n'
                f'最低价: {data_with_indicators["low"].min():.2f}\n'
                f'平均价: {data_with_indicators["close"].mean():.2f}'
            )
            ax1.text(0.02, 0.98, price_stats, transform=ax1.transAxes, verticalalignment='top',
                     bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8),
                     fontsize=9, linespacing=1.5)  # 字体从11减小到9

            # 副图：成交量
            bar_colors = ['#FF9999' if data_with_indicators['close'].iloc[i] >= data_with_indicators['open'].iloc[
                i] else '#99FF99'
                          for i in range(len(data_with_indicators))]
            ax2.bar(data_with_indicators.index, data_with_indicators['volume'], color=bar_colors, alpha=0.7,
                    label='成交量')
            ax2.set_title('成交量', fontsize=12, fontweight='bold', pad=10)  # 字体从16减小到12
            ax2.set_xlabel('日期', fontsize=10)  # 字体从14减小到10
            ax2.set_ylabel('成交量', fontsize=10)  # 字体从14减小到10
            ax2.legend(fontsize=8)  # 字体从12减小到8
            ax2.grid(True, alpha=0.3)

            # 添加成交量统计信息 - 减小字体
            volume_stats = (
                f'最大成交量: {data_with_indicators["volume"].max():,.0f}\n'
                f'平均成交量: {data_with_indicators["volume"].mean():,.0f}\n'
                f'最新成交量: {data_with_indicators["volume"].iloc[-1]:,.0f}'
            )
            ax2.text(0.02, 0.98, volume_stats, transform=ax2.transAxes, verticalalignment='top',
                     bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8),
                     fontsize=9, linespacing=1.5)  # 字体从11减小到9

            plt.tight_layout(pad=2.0)  # 从3.0减小到2.0

            # 保存为JPG - 降低DPI和质量
            jpg_filename = os.path.join(stock_dir, f'{code}_{freq}_kline.jpg')
            plt.savefig(jpg_filename, dpi=150, bbox_inches='tight', format='jpg')  # DPI从300降低到150，添加压缩优化，质量85%
            jpg_files.append(jpg_filename)

            # 获取文件大小
            file_size = os.path.getsize(jpg_filename) / 1024  # KB
            print(f"{freq_name}K线图表已保存为: {jpg_filename} (大小: {file_size:.1f} KB)")
            plt.close()  # 关闭图表以释放内存

            # 保存图片路径到数据库
            save_image_path_to_db(code, 'individual', freq, jpg_filename)

    return jpg_files


def create_monthly_candle_chart_jpg(data_dict, code, stock_name="", save_dir="gupiaoimagenew"):
    """创建月K线蜡烛图并保存为JPG - 优化版本"""
    # 创建股票专属目录
    stock_dir = os.path.join(save_dir, code)
    if not os.path.exists(stock_dir):
        os.makedirs(stock_dir)

    jpg_files = []

    # 只处理月K线数据
    data = data_dict.get('monthly')
    if data is not None and len(data) > 0:
        # 计算技术指标
        data_with_indicators = calculate_technical_indicators(data.copy())

        # 减小图片尺寸
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))  # 从 (16, 12) 减小到 (12, 8)

        # 创建K线图效果 - 减小线宽
        for i, (idx, row) in enumerate(data_with_indicators.iterrows()):
            color = 'red' if row['close'] >= row['open'] else 'green'
            # 绘制上下影线 - 减小线宽
            ax1.plot([idx, idx], [row['low'], row['high']], color=color, linewidth=0.8)  # 线宽从1减小到0.8
            # 绘制实体 - 减小线宽
            ax1.plot([idx, idx], [min(row['open'], row['close']), max(row['open'], row['close'])],
                     color=color, linewidth=4)  # 线宽从6减小到4

        ax1.set_title(f'{stock_name} 月K线图（蜡烛图）', fontsize=14, fontweight='bold', pad=15)  # 字体从18减小到14
        ax1.set_ylabel('价格 (元)', fontsize=12)  # 字体从14减小到12
        ax1.grid(True, alpha=0.3)

        # 添加移动平均线 - 减小线宽
        if 'MA5' in data_with_indicators.columns:
            ax1.plot(data_with_indicators.index, data_with_indicators['MA5'],
                     '--', color='orange', alpha=0.8, label='MA5', linewidth=1.5)  # 线宽从2减小到1.5
        if 'MA10' in data_with_indicators.columns:
            ax1.plot(data_with_indicators.index, data_with_indicators['MA10'],
                     '--', color='purple', alpha=0.8, label='MA10', linewidth=1.5)  # 线宽从2减小到1.5
        ax1.legend(fontsize=10)  # 字体从12减小到10

        # 成交量（颜色与涨跌一致）
        volume_colors = [
            'red' if data_with_indicators['close'].iloc[i] >= data_with_indicators['open'].iloc[i] else 'green'
            for i in range(len(data_with_indicators))]
        ax2.bar(data_with_indicators.index, data_with_indicators['volume'], color=volume_colors, alpha=0.7)
        ax2.set_title('成交量', fontsize=12, fontweight='bold', pad=10)  # 字体从16减小到12
        ax2.set_xlabel('日期', fontsize=10)  # 字体从14减小到10
        ax2.set_ylabel('成交量', fontsize=10)  # 字体从14减小到10
        ax2.grid(True, alpha=0.3)

        plt.tight_layout(pad=2.0)  # 从3.0减小到2.0

        # 保存为JPG - 降低DPI和质量
        jpg_filename = os.path.join(stock_dir, f'{code}_monthly_candle_chart.jpg')
        plt.savefig(jpg_filename, dpi=150, bbox_inches='tight', format='jpg')  # DPI从300降低到150，添加压缩优化，质量85%
        jpg_files.append(jpg_filename)

        # 获取文件大小
        file_size = os.path.getsize(jpg_filename) / 1024  # KB
        print(f"月K线蜡烛图已保存为: {jpg_filename} (大小: {file_size:.1f} KB)")
        plt.close()

        # 保存图片路径到数据库
        save_image_path_to_db(code, 'candle', 'monthly', jpg_filename)

    return jpg_files


def process_single_stock(code, name, start_date, end_date, save_directory):
    """处理单只股票的数据和图表生成"""
    print(f"\n{'=' * 50}")  # 减小分隔符长度
    print(f"开始处理股票: {code} - {name}")
    print(f"{'=' * 50}")

    try:
        # 获取所有周期的K线数据
        all_data = get_all_kline_data(code, start_date, end_date)

        # 检查数据是否获取成功
        if all_data is None:
            print(f"股票 {code} 数据获取失败，跳过")
            return False

        success_count = sum(1 for data in all_data.values() if data is not None and len(data) > 0)
        if success_count == 0:
            print(f"股票 {code} 所有数据获取失败，跳过")
            return False

        # 打印数据概览
        for freq, data in all_data.items():
            if data is not None and len(data) > 0:
                freq_name = {'monthly': '月', 'quarterly': '季', 'yearly': '年'}.get(freq, freq)
                print(f"{freq_name}K线数据: {len(data)} 个数据点")

        # 生成需要的JPG图表
        comparison_chart = create_comparison_chart_jpg(all_data, code, name, save_directory)
        individual_charts = create_individual_chart_jpg(all_data, code, name, save_directory)  # 年K线和季K线
        monthly_candle_chart = create_monthly_candle_chart_jpg(all_data, code, name, save_directory)  # 月蜡烛图

        total_jpg_files = 1 + len(individual_charts) + len(monthly_candle_chart)
        print(f"✓ 成功生成 {total_jpg_files} 个图表文件")
        return True

    except Exception as e:
        print(f"处理股票 {code} 时出错: {e}")
        import traceback
        traceback.print_exc()
        return False


# 主程序
if __name__ == "__main__":
    # 设置参数k线数据
    start_date = "2010-01-01"
    end_date = "2025-10-31"
    save_directory = "gupiaoimagenew"  # 图片保存目录

    # 创建主目录
    if not os.path.exists(save_directory):
        os.makedirs(save_directory)

    # 登录baostock（全局一次）
    print("登录Baostock系统...")
    lg = bs.login()
    if lg.error_code != '0':
        print(f'Baostock登录失败: {lg.error_msg}')
        exit()

    try:
        # 获取所有股票列表
        print("从数据库获取股票列表...")
        all_stocks = get_all_stocks_from_db()

        if not all_stocks:
            print("未获取到任何股票数据")
            exit()

        print(f"共获取到 {len(all_stocks)} 只股票")

        # 统计变量
        success_count = 0
        failed_count = 0
        processed_count = 0

        # 处理每只股票
        for code, name in all_stocks:
            processed_count += 1
            print(f"\n进度: {processed_count}/{len(all_stocks)}")

            if process_single_stock(code, name, start_date, end_date, save_directory):
                success_count += 1
            else:
                failed_count += 1

            # 添加延迟，避免请求过于频繁
            time.sleep(1)

        # 输出总结报告
        print(f"\n{'=' * 50}")
        print("批量处理完成总结")
        print(f"{'=' * 50}")
        print(f"总处理股票数: {len(all_stocks)}")
        print(f"成功生成图表: {success_count}")
        print(f"处理失败: {failed_count}")
        print(f"成功率: {success_count / len(all_stocks) * 100:.1f}%")
        print(f"所有图表文件已保存到: {save_directory} 目录")

    except Exception as e:
        print(f"程序执行出错: {e}")
        import traceback

        traceback.print_exc()

    finally:
        # 确保登出
        try:
            bs.logout()
            print("已登出系统")
        except:
            pass